package com.bw;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.dataproc.SplitBatchOp;
import com.alibaba.alink.operator.batch.source.CsvSourceBatchOp;
import com.alibaba.alink.params.feature.HasEncodeWithoutWoe;
import com.alibaba.alink.pipeline.PipelineModel;
import com.alibaba.alink.pipeline.dataproc.vector.VectorAssembler;
import com.alibaba.alink.pipeline.feature.OneHotEncoder;
import com.alibaba.alink.pipeline.regression.DecisionTreeRegressionModel;
import com.alibaba.alink.pipeline.regression.DecisionTreeRegressor;
import com.alibaba.alink.pipeline.regression.LinearRegression;
import com.alibaba.alink.pipeline.regression.LinearRegressionModel;
import com.alibaba.alink.pipeline.tuning.GridSearchCV;
import com.alibaba.alink.pipeline.tuning.GridSearchCVModel;
import com.alibaba.alink.pipeline.tuning.ParamGrid;
import com.alibaba.alink.pipeline.tuning.RegressionTuningEvaluator;

public class Test2 {
    public static void main(String[] args) throws Exception {
        BatchOperator.setParallelism(1);
        String schemaStr = "car_ID int,symboling int,CarName string,fueltype string,aspiration string,doornumber string,carbody string,drivewheel string,enginelocation string,wheelbase double,carlength double,carwidth double,carheight double,curbweight int,enginetype string,cylindernumber string,enginesize int,fuelsystem string,boreratio double,stroke double,compressionratio double,horsepower int,peakrpm int,citympg int,highwaympg int,price double";

        //CarName  fueltype  aspiration doornumber carbody drivewheel
        CsvSourceBatchOp source = new CsvSourceBatchOp()
                .setFilePath("datafile/test2.csv")
                .setSchemaStr(schemaStr)
                .setFieldDelimiter(",")
                .setIgnoreFirstLine(true);
//        source.print();

        //编码
        OneHotEncoder one_hot_06 = new OneHotEncoder().setEncode(HasEncodeWithoutWoe.Encode.INDEX)
                .setSelectedCols("CarName", "fueltype", "aspiration", "doornumber", "carbody", "drivewheel")
//                .setOutputCols("vec");
                .setOutputCols("CarName_new", "fueltype_new", "aspiration_new", "doornumber_new", "carbody_new", "drivewheel_new");

        BatchOperator<?> source_06 = one_hot_06.fit(source).transform(source);
        source_06.print();


        // 选择特征
        String[] features = new String[]{"symboling", "CarName_new", "fueltype_new", "aspiration_new",
                "doornumber_new", "carbody_new", "drivewheel_new","highwaympg"};

        // 目标列
        String lable = "price";



        // 向量聚合
//        VectorAssembler res = new VectorAssembler()
//                .setSelectedCols(features)
//                .setOutputCol("vec");
//        BatchOperator<?> VectorResult =res.transform(source_06);
//        VectorResult.print();

        // 拆分数据集
        BatchOperator<?> trainData = new SplitBatchOp().setFraction(0.8).linkFrom(source_06);
        BatchOperator<?> testData = trainData.getSideOutput(0);

        testData.print();
        trainData.print();

        // 线性回归
        LinearRegression lr = new LinearRegression()
                .setFeatureCols(features)
                .setL1(0.1)
                .setLabelCol(lable)
                .setPredictionCol("pred")
                .enableLazyPrintModelInfo("ModelInfo");
        // 训练模型
        LinearRegressionModel lr_model = lr.fit(trainData);
        BatchOperator<?> lr_result=lr_model.transform(testData);
        lr_result.print();


        // 用决策树回归
        DecisionTreeRegressor dt = new DecisionTreeRegressor()
                .setPredictionCol("pred")
                .setNumThreads(10)
                .setLabelCol(lable)
                .setFeatureCols(features);
        BatchOperator<?> dt_result = dt
                .fit(trainData)
                .transform(testData);
        dt_result.print();




        RegressionTuningEvaluator eval = new RegressionTuningEvaluator()
                .setLabelCol(lable)
                .setPredictionCol("pred")
                .setTuningRegressionMetric("RMSE");
        // 线下回归调优

       ParamGrid lr_paramGrid = new ParamGrid()
                    .addGrid(lr, LinearRegression.L_1, new Double[]{0.1, 0.2});
        GridSearchCV lr_cv = new GridSearchCV()
                .setEstimator(lr)
                .setParamGrid(lr_paramGrid)
                .setTuningEvaluator(eval)
                .setNumFolds(2)
                .enableLazyPrintTrainInfo("TrainInfo");

        GridSearchCVModel lr_model1 = lr_cv.fit(trainData);
        PipelineModel lr_best_model = lr_model1.getBestPipelineModel();
        BatchOperator<?> lr_test_result = lr_best_model.transform(testData);
        lr_best_model.save("datafile/yk02_bast_lr_model",true);






        // 决策树调优
        ParamGrid tree_paramGrid = new ParamGrid()
                .addGrid(dt, DecisionTreeRegressor.MIN_SAMPLES_PER_LEAF, new Integer[]{1, 3});
        GridSearchCV tree_cv = new GridSearchCV()
                .setEstimator(dt)
                .setParamGrid(tree_paramGrid)
                .setTuningEvaluator(eval)
                .setNumFolds(2)
                .enableLazyPrintTrainInfo("TrainInfo");
        GridSearchCVModel tree_model = tree_cv.fit(trainData);
        PipelineModel tree_best_model = tree_model.getBestPipelineModel();
//        BatchOperator<?> lr_test_result = tree_best_model.transform(testData);
        tree_best_model.save("datafile/test_alink_file_sink",true);



        BatchOperator.execute();








    }
}
